import json
import os
import time
from typing import Dict, List, Any

from knowledge_editor_model import KnowledgeEditorFactory
from evaluation import Evaluator, calculate_overall_score

# 定义数据集路径
WIKI_RECENT_TRAIN = 'data/Wiki_recent/recent_train.json'
WIKI_RECENT_TEST = 'data/Wiki_recent/recent_test.json'
ZSRE_TEST = 'data/ZsRE/ZsRE-test-all.json'
SAFEEDIT_TRAIN = 'data/SafeEdit/SafeEdit_train.json'
SAFEEDIT_TEST = 'data/SafeEdit/SafeEdit_test_ccks.json'

# 定义结果文件路径
RESULT_FILE = 'team_name_result.json'

# 加载数据集
def load_data(file_path: str) -> List[Dict]:
    """加载数据集"""
    with open(file_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    return data

# 处理Wiki_recent数据集（知识新增）
def process_wiki_recent(test_data: List[Dict], sample_limit: int = None) -> List[Dict]:
    """处理Wiki_recent数据集"""
    print("\n处理Wiki_recent数据集（知识新增）...")
    results = []
    
    # 创建Wiki_recent编辑器
    wiki_editor = KnowledgeEditorFactory.create_editor('wiki_recent')
    evaluator = Evaluator()
    
    # 限制处理的样本数量
    if sample_limit:
        test_data = test_data[:sample_limit]
    
    for idx, item in enumerate(test_data):
        print(f"处理Wiki_recent样本 {idx+1}/{len(test_data)}")
        
        # 编辑前评估
        pre_answer = wiki_editor.answer_question(item['prompt'])
        pre_metrics = {
            'answer': pre_answer,
            'success': pre_answer == item['target_new']
        }
        
        # 执行知识编辑
        edited_result = wiki_editor.edit_knowledge(item)
        
        # 编辑后评估
        post_answer = wiki_editor.answer_question(item['prompt'])
        post_metrics = {
            'answer': post_answer,
            'success': post_answer == item['target_new'],
            'portability_answers': edited_result.get('portability_answers', {}),
            'locality_answers': edited_result.get('locality_answers', {})
        }
        
        # 计算评估指标
        evaluation = evaluator.evaluate_wiki_recent(pre_metrics, post_metrics, item)
        
        # 构建结果
        result = {
            'pre': pre_metrics,
            'case_id': item['subject'],
            'requested_rewrite': {
                'subject': item['subject'],
                'prompt': item['prompt'],
                'target_new': item['target_new']
            },
            'time': edited_result['edit_time'],
            'post': post_metrics,
            'es': evaluation['es'],
            'port': evaluation['port'],
            'loc': evaluation['loc'],
            'flue': evaluation['flue']
        }
        
        results.append(result)
    
    return results

# 处理ZsRE数据集（知识修改）
def process_zsre(test_data: List[Dict], sample_limit: int = None) -> List[Dict]:
    """处理ZsRE数据集"""
    print("\n处理ZsRE数据集（知识修改）...")
    results = []
    
    # 创建ZsRE编辑器
    zsre_editor = KnowledgeEditorFactory.create_editor('zsre')
    evaluator = Evaluator()
    
    # 限制处理的样本数量
    if sample_limit:
        test_data = test_data[:sample_limit]
    
    for idx, item in enumerate(test_data):
        print(f"处理ZsRE样本 {idx+1}/{len(test_data)}")
        
        # 编辑前评估
        pre_answer = zsre_editor.answer_question(item['prompt'])
        pre_metrics = {
            'answer': pre_answer,
            'success': False  # 编辑前应该是错误的答案
        }
        
        # 执行知识编辑
        edited_result = zsre_editor.edit_knowledge(item)
        
        # 编辑后评估
        post_answer = zsre_editor.answer_question(item['prompt'])
        post_metrics = {
            'answer': post_answer,
            'success': post_answer == item['target_new'],
            'portability_answers': edited_result.get('portability_answers', {}),
            'locality_answers': edited_result.get('locality_answers', {})
        }
        
        # 计算评估指标
        evaluation = evaluator.evaluate_zsre(pre_metrics, post_metrics, item)
        
        # 构建结果
        result = {
            'pre': pre_metrics,
            'case_id': item['subject'],
            'requested_rewrite': {
                'subject': item['subject'],
                'prompt': item['prompt'],
                'target_new': item['target_new']
            },
            'time': edited_result['edit_time'],
            'post': post_metrics,
            'es': evaluation['es'],
            'port': evaluation['port'],
            'loc': evaluation['loc'],
            'flue': evaluation['flue']
        }
        
        results.append(result)
    
    return results

# 处理SafeEdit数据集（知识删除）
def process_safeedit(test_data: List[Dict], sample_limit: int = None) -> List[Dict]:
    """处理SafeEdit数据集"""
    print("\n处理SafeEdit数据集（知识删除）...")
    results = []
    
    # 创建SafeEdit编辑器
    safeedit_editor = KnowledgeEditorFactory.create_editor('safeedit')
    evaluator = Evaluator()
    
    # 限制处理的样本数量
    if sample_limit:
        test_data = test_data[:sample_limit]
    
    for idx, item in enumerate(test_data):
        print(f"处理SafeEdit样本 {idx+1}/{len(test_data)}")
        
        # 编辑前评估
        pre_answer = safeedit_editor.answer_question(item['question'])
        pre_metrics = {
            'answer': pre_answer,
            'success': False  # 编辑前应该是不安全的答案
        }
        
        # 执行知识编辑
        edited_result = safeedit_editor.edit_knowledge(item)
        
        # 编辑后评估
        post_answer = safeedit_editor.answer_question(item['question'])
        post_metrics = {
            'answer': post_answer,
            'success': True,  # 编辑后应该是安全的答案
            'generalization_answers': edited_result.get('generalization_answers', {}),
            'constrain_answer': edited_result.get('constrain_answer', '')
        }
        
        # 计算评估指标
        evaluation = evaluator.evaluate_safeedit(pre_metrics, post_metrics, item)
        
        # 构建结果
        result = {
            'pre': pre_metrics,
            'case_id': item['id'],
            'requested_rewrite': {
                'id': item['id'],
                'category': item['unsafety category'],
                'question': item['question'],
                'safe_generation': item['safe generation']
            },
            'time': edited_result['edit_time'],
            'post': post_metrics,
            'es': evaluation['es'],
            'port': evaluation['port'],
            'loc': evaluation['loc'],
            'flue': evaluation['flue']
        }
        
        results.append(result)
    
    return results

# 合并结果并保存
def merge_and_save_results(wiki_results: List[Dict], zsre_results: List[Dict], safeedit_results: List[Dict], output_file: str) -> None:
    """合并结果并保存"""
    # 合并所有结果
    all_results = wiki_results + zsre_results + safeedit_results
    
    # 计算总体评分
    overall_scores = calculate_overall_score(all_results)
    print("\n总体评分:")
    print(f"Wiki_recent: ES={overall_scores['wiki']['es']:.4f}, PORT={overall_scores['wiki']['port']:.4f}, LOC={overall_scores['wiki']['loc']:.4f}, FLUE={overall_scores['wiki']['flue']:.4f}, Total={overall_scores['wiki']['total']:.4f}")
    print(f"ZsRE: ES={overall_scores['zsre']['es']:.4f}, PORT={overall_scores['zsre']['port']:.4f}, LOC={overall_scores['zsre']['loc']:.4f}, FLUE={overall_scores['zsre']['flue']:.4f}, Total={overall_scores['zsre']['total']:.4f}")
    print(f"SafeEdit: ES={overall_scores['safeedit']['es']:.4f}, PORT={overall_scores['safeedit']['port']:.4f}, LOC={overall_scores['safeedit']['loc']:.4f}, FLUE={overall_scores['safeedit']['flue']:.4f}, Total={overall_scores['safeedit']['total']:.4f}")
    print(f"Overall: ES={overall_scores['overall']['es']:.4f}, PORT={overall_scores['overall']['port']:.4f}, LOC={overall_scores['overall']['loc']:.4f}, FLUE={overall_scores['overall']['flue']:.4f}, Total={overall_scores['overall']['total']:.4f}")
    
    # 准备提交结果
    submission_results = []
    for result in all_results:
        submission_result = {
            'pre': result['pre'],
            'case_id': result['case_id'],
            'requested_rewrite': result['requested_rewrite'],
            'time': result['time'],
            'post': result['post']
        }
        submission_results.append(submission_result)
    
    # 保存结果
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(submission_results, f, ensure_ascii=False, indent=4)
    
    print(f"\n结果已保存到 {output_file}")

# 主函数
def main():
    print("开始知识编辑任务...")
    
    # 加载数据集
    print("加载数据集...")
    wiki_test_data = load_data(WIKI_RECENT_TEST)
    zsre_test_data = load_data(ZSRE_TEST)
    safeedit_test_data = load_data(SAFEEDIT_TEST)
    
    # 设置样本限制（用于测试，实际提交时可以设置为None处理全部数据）
    sample_limit = 5  # 每个数据集只处理前5个样本
    
    # 处理各个数据集
    wiki_results = process_wiki_recent(wiki_test_data, sample_limit)
    zsre_results = process_zsre(zsre_test_data, sample_limit)
    safeedit_results = process_safeedit(safeedit_test_data, sample_limit)
    
    # 合并结果并保存
    merge_and_save_results(wiki_results, zsre_results, safeedit_results, RESULT_FILE)
    
    print("知识编辑任务完成!")

if __name__ == "__main__":
    main()